#!/bin/bash
# This script creates a SquashFS image with the kernel modules and initramfs image.
# Author: crims0n <https://minios.dev>
set -euo pipefail

# Display help information
show_help() {
    cat <<EOF
Usage: $0 [OPTIONS]

Options:
  --help        Show this help message and exit

This script creates a SquashFS image with the kernel modules and an initramfs image.
EOF
    exit 0
}

# Check for --help option
if [[ "${1-}" == "--help" ]]; then
    show_help
fi

# Function to print steps
print_step() {
    echo -e "\033[1;32m* $1...\033[0m"
}

# Check for root privileges
if [[ $EUID -ne 0 ]]; then
    echo "This script must be run as root."
    exit 1
fi

# Function to select the kernel version
select_kernel() {
    print_step "Selecting kernel version"
    # Get the list of installed kernels
    mapfile -t kernels < <(ls /lib/modules)
    echo "  Installed kernels found: ${kernels[*]}"
    # Use whiptail to select the kernel, even if there is only one
    options=()
    for kernel in "${kernels[@]}"; do
        options+=("$kernel" "")
    done
    kernel=$(whiptail --title "Select the Kernel" --menu "Choose a kernel:" 15 60 4 "${options[@]}" 3>&1 1>&2 2>&3)
    if [ -z "$kernel" ]; then
        echo "No kernel selected. Exiting."
        exit 1
    fi
    echo "  Selected kernel: $kernel"
}

# Declare variables
temp_dir=$(mktemp -d)
sqfs_mod=""
vmlinuz_path=""
kernel=""

# Function to clean up on exit
cleanup() {
    print_step "Cleaning up temporary files and unmounting directories"
    cp "${temp_dir}/livekit/initramfs.log" initramfs-$kernel.log
    rm -rf "$temp_dir"
    echo "  Cleanup complete."
}
trap cleanup EXIT

# Select kernel version
select_kernel

# Determine the path to the vmlinuz file
if [ -f "/boot/vmlinuz-$kernel" ]; then
    vmlinuz_path="/boot/vmlinuz-$kernel"
elif [ -f "/run/initramfs/memory/data/minios/boot/vmlinuz-$kernel" ]; then
    vmlinuz_path="/run/initramfs/memory/data/minios/boot/vmlinuz-$kernel"
elif [ -f "/run/initramfs/memory/data/minios/boot/vmlinuz" ]; then
    vmlinuz_path="/run/initramfs/memory/data/minios/boot/vmlinuz"
fi
if [ ! -f "$vmlinuz_path" ]; then
    echo "Error: The vmlinuz file for the selected kernel does not exist."
    exit 1
fi

# Copy the kernel image
print_step "Copying kernel image"
cp "$vmlinuz_path" "./vmlinuz-$kernel"
echo "  Kernel image copied to: ./vmlinuz-$kernel"

# Determine the modules directory
print_step "Determining modules directory"
if [ -d "/lib/modules/$kernel" ]; then
    sqfs_mod="/lib/modules"
else
    echo "Error: The modules directory does not exist."
    exit 1
fi

# Create SquashFS file for kernel modules
print_step "Creating SquashFS file for kernel modules"
mkdir -p "${temp_dir}/kernel/$sqfs_mod"
cp -r "/lib/modules/$kernel" "${temp_dir}/kernel/$sqfs_mod"

if [ -f "01-kernel-$kernel.sb" ]; then
    rm "01-kernel-$kernel.sb"
fi
mksquashfs "${temp_dir}/kernel" "01-kernel-$kernel.sb" -comp zstd -Xcompression-level 19 -b 1024K -always-use-fragments -noappend -quiet -progress
rm -rf "${temp_dir}/kernel"
echo "  SquashFS file created: 01-kernel-$kernel.sb"

# Copy necessary binaries and libraries
print_step "Copying necessary binaries and libraries"
mkdir -p "$temp_dir/livekit/"{static,scripts,other}

for bin in busybox eject resize2fs mc ncurses-menu blkid; do
    if [ -f "/run/initramfs/bin/$bin" ]; then
        cp "/run/initramfs/bin/$bin" "$temp_dir/livekit/static"
    fi
done
for bin in mount.dynfilefs mount.httpfs2 mount.ntfs-3g; do
    if [ -f "/run/initramfs/bin/@$bin" ]; then
        cp "/run/initramfs/bin/@$bin" "$temp_dir/livekit/static/$bin"
    fi
done
for script in minios-boot; do
    if [ -f "/run/initramfs/bin/$script" ]; then
        cp "/run/initramfs/bin/$script" "$temp_dir/livekit/scripts"
    fi
done

for file in config livekitlib; do
    if [ -f "/run/initramfs/lib/$file" ]; then
        cp "/run/initramfs/lib/$file" "$temp_dir/livekit/$file"
    fi
done
for file in init shutdown; do
    if [ -f "/run/initramfs/$file" ]; then
        cp "/run/initramfs/$file" "$temp_dir/livekit/$file"
    fi
done

if [ -f "/run/initramfs/usr/share/terminfo/l/linux" ]; then
    cp "/run/initramfs/usr/share/terminfo/l/linux" "$temp_dir/livekit/other/linux3.0"
fi

# Create initramfs creation script
cat <<'EOF' >"$temp_dir/livekit/initramfs_create"
#!/bin/bash
# Create initramfs image
# Author: Tomas M <http://www.linux-live.org/>
# Author: crims0n <https://minios.dev>
set -euo pipefail

# Check if KERNEL is supplied as an argument
[ ! "${1-}" ] && echo "Builds LiveKit initrfs.img
    Usage:   $0 [\$KERNEL]
    Example: $0 6.8.11-mos-amd64" && exit 1

# Assign kernel version from the first script argument
KERNEL="$1"

# Set config file path
if [ -f "../../container_config" ]; then
   CONFIG="../../container_config"
elif [ -f "../../config" ]; then
   CONFIG="../../config"
elif [ -f "./config" ]; then
   CONFIG="./config"
elif [ -f "/etc/minios/config" ]; then
   CONFIG="/etc/minios/config"
fi

# Assign config values
LIVEKITNAME="minios"
NETWORK="true"
CLOUD="false"

# Set paths and filename
LMK="lib/modules/$KERNEL"
INITRAMFS=/tmp/$LIVEKITNAME-initramfs-$$
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
LOG_FILE="$SCRIPT_DIR/initramfs.log"

# copy file to initramfs tree, including
# all library dependencies (as shown by ldd)
# $1 = file to copy (full path)
copy_including_deps() {
   # if source doesn't exist or target exists, do nothing
   if [ ! -e "$1" -o -e "$INITRAMFS"/"$1" ]; then
      return
   fi

   cp -R --parents "$1" "$INITRAMFS"
   if [ -L "$1" ]; then
      DIR="$(dirname "$1")"
      LNK="$(readlink "$1")"
      copy_including_deps "$(
         cd "$DIR"
         realpath -s "$LNK"
      )"
   fi

   # Use `ldd` only if $1 is a file and is not a kernel module
   if [ -f "$1" ] && [[ "$1" != "/lib/modules/"* ]]; then
      ldd "$1" 2>/dev/null | sed -r "s/.*=>|[(].*//g" | sed -r "s:^\\s+|\\s+$::g" |
         while read LIB; do
            copy_including_deps "$LIB"
         done
   fi

   for MOD in $(find "$1" -type f | grep .ko); do
      for DEP in $(cat /$LMK/modules.dep | fgrep "/$(basename $MOD):"); do
         copy_including_deps "/$LMK/$DEP"
      done
   done

   shift
   if [ "${1-}" != "" ]; then
      copy_including_deps "$@"
   fi
}

# Redirect trace to log file and start tracing commands
exec 19>"$LOG_FILE"
BASH_XTRACEFD=19
set -x

if [ -d $INITRAMFS ]; then
   rm -Rf $INITRAMFS
fi
mkdir -p $INITRAMFS/{bin,dev,etc,lib,lib64,mnt,proc,root,run,sys,tmp,usr,var/log}
ln -s bin $INITRAMFS/sbin

cp static/busybox $INITRAMFS/bin
cp static/eject $INITRAMFS/bin
cp static/resize2fs $INITRAMFS/bin
cp static/mc $INITRAMFS/bin
cp static/mount.dynfilefs $INITRAMFS/bin/@mount.dynfilefs
cp static/mount.httpfs2 $INITRAMFS/bin/@mount.httpfs2
cp static/ncurses-menu $INITRAMFS/bin
mkdir -p $INITRAMFS/usr/share/terminfo/l
cp other/linux3.0 $INITRAMFS/usr/share/terminfo/l/linux

# if the kernel has a version lower than 5.15, it does not support NTFS3, we use NTFS-3G.
if [[ "$(printf '%s\n' "$KERNEL" 5.15 | sort -Vr | head -n1)" == 5.15 ]]; then
   cp static/mount.ntfs-3g $INITRAMFS/bin/@mount.ntfs-3g
fi

cp static/blkid $INITRAMFS/bin
cp scripts/minios-boot $INITRAMFS/bin
chmod a+x $INITRAMFS/bin/*
mkdir -p $INITRAMFS/etc/modprobe.d
echo "options loop max_loop=64" >$INITRAMFS/etc/modprobe.d/local-loop.conf

$INITRAMFS/bin/busybox | grep , | grep -v Copyright | tr "," " " | while read LINE; do
   for TOOL in $LINE; do
      if [ ! -e $INITRAMFS/bin/$TOOL ]; then
         ln -s busybox $INITRAMFS/bin/$TOOL
      fi
   done
done
rm -f $INITRAMFS/{s,}bin/init

mknod $INITRAMFS/dev/console c 5 1
mknod $INITRAMFS/dev/null c 1 3
mknod $INITRAMFS/dev/ram0 b 1 0
mknod $INITRAMFS/dev/tty1 c 4 1
mknod $INITRAMFS/dev/tty2 c 4 2
mknod $INITRAMFS/dev/tty3 c 4 3
mknod $INITRAMFS/dev/tty4 c 4 4

copy_including_deps /$LMK/kernel/fs/aufs
copy_including_deps /$LMK/kernel/fs/overlayfs
copy_including_deps /$LMK/kernel/fs/ext2
copy_including_deps /$LMK/kernel/fs/ext3
copy_including_deps /$LMK/kernel/fs/ext4
copy_including_deps /$LMK/kernel/fs/exfat
copy_including_deps /$LMK/kernel/fs/fat
copy_including_deps /$LMK/kernel/fs/nls
copy_including_deps /$LMK/kernel/fs/fuse
copy_including_deps /$LMK/kernel/fs/isofs
copy_including_deps /$LMK/kernel/fs/ntfs
copy_including_deps /$LMK/kernel/fs/ntfs3
copy_including_deps /$LMK/kernel/fs/squashfs
copy_including_deps /$LMK/kernel/fs/btrfs
copy_including_deps /$LMK/kernel/fs/xfs
copy_including_deps /$LMK/kernel/fs/efivarfs # for Clonezilla

copy_including_deps /$LMK/kernel/crypto/lz4*.*
copy_including_deps /$LMK/kernel/crypto/zstd.*

# crc32c is needed for ext4, but I don't know which one, add them all, they are small
find /$LMK/kernel/ | grep crc32c | while read LINE; do
   copy_including_deps $LINE
done

copy_including_deps /$LMK/kernel/drivers/staging/zsmalloc # needed by zram
copy_including_deps /$LMK/kernel/drivers/block/zram
copy_including_deps /$LMK/kernel/drivers/block/loop.*
copy_including_deps /$LMK/kernel/drivers/block/nbd.* # for Rescuezilla
copy_including_deps /$LMK/kernel/drivers/md/dm-mod.* # for Ventoy

# usb drivers
copy_including_deps /$LMK/kernel/drivers/usb/storage
copy_including_deps /$LMK/kernel/drivers/usb/host
copy_including_deps /$LMK/kernel/drivers/usb/common
copy_including_deps /$LMK/kernel/drivers/usb/core
copy_including_deps /$LMK/kernel/drivers/hid/usbhid
copy_including_deps /$LMK/kernel/drivers/hid/hid.*
copy_including_deps /$LMK/kernel/drivers/hid/uhid.*
copy_including_deps /$LMK/kernel/drivers/hid/hid-generic.*

# disk and cdrom drivers
copy_including_deps /$LMK/kernel/drivers/cdrom
copy_including_deps /$LMK/kernel/drivers/scsi/sr_mod.*
copy_including_deps /$LMK/kernel/drivers/scsi/sd_mod.*
copy_including_deps /$LMK/kernel/drivers/scsi/scsi_mod.*
copy_including_deps /$LMK/kernel/drivers/scsi/sg.*
# virtual disk drivers
copy_including_deps /$LMK/kernel/drivers/scsi/hv_storvsc.*

if [ "$CLOUD" = "true" ]; then
   copy_including_deps /$LMK/kernel/drivers/virtio/virtio.*
   copy_including_deps /$LMK/kernel/drivers/virtio/virtio_mmio.*
   copy_including_deps /$LMK/kernel/drivers/virtio/virtio_pci.*
   copy_including_deps /$LMK/kernel/drivers/virtio/virtio_ring.*
   copy_including_deps /$LMK/kernel/drivers/scsi/vmw_pvscsi.*
   copy_including_deps /$LMK/kernel/drivers/scsi/virtio_scsi.*
   copy_including_deps /$LMK/kernel/drivers/block/virtio_blk.*
   copy_including_deps /$LMK/kernel/drivers/scsi/hv_storvsc.*
fi

copy_including_deps /$LMK/kernel/drivers/ata
copy_including_deps /$LMK/kernel/drivers/nvme
copy_including_deps /$LMK/kernel/drivers/mmc

# network support drivers
if [ "$NETWORK" = "true" ]; then
   copy_including_deps /$LMK/kernel/drivers/net/ethernet
   copy_including_deps /$LMK/kernel/drivers/net/phy
   if [ "$CLOUD" = "true" ]; then
      copy_including_deps /$LMK/kernel/drivers/net/vmxnet3
      copy_including_deps /$LMK/kernel/drivers/net/virtio_net.*
   fi
fi

# copy all custom-built modules
copy_including_deps /$LMK/updates

copy_including_deps /$LMK/modules.*

copy_including_deps /usr/share/terminfo/l/linux

find $INITRAMFS -name "*.ko.gz" -exec gunzip {} \;
find $INITRAMFS -name "*.ko.xz" -exec unxz {} \;

# trim modules.order file. Perhaps we could remove it entirely
MODULEORDER="$(
   cd "$INITRAMFS/$LMK/"
   find -name "*.ko" | sed -r "s:^./::g" | tr "\n" "|" | sed -r "s:[.]:.:g"
)"
cat $INITRAMFS/$LMK/modules.order | sed -r "s/.ko.gz\$/.ko/" | grep -E "$MODULEORDER"/foo/bar >$INITRAMFS/$LMK/_
mv $INITRAMFS/$LMK/_ $INITRAMFS/$LMK/modules.order

depmod -b $INITRAMFS $KERNEL

echo "root::0:0::/root:/bin/bash" >$INITRAMFS/etc/passwd
touch $INITRAMFS/etc/{m,fs}tab

cp init $INITRAMFS
chmod a+x $INITRAMFS/init
cp shutdown $INITRAMFS
chmod a+x $INITRAMFS/shutdown
ln -s ../init $INITRAMFS/bin/init
cp ./livekitlib $INITRAMFS/lib/
cp $CONFIG $INITRAMFS/lib/config

cd $INITRAMFS
find . -print | cpio -o -H newc 2>/dev/null | xz -T0 -f --extreme --check=crc32 >$INITRAMFS.img
echo $INITRAMFS.img

cd ..
rm -Rf $INITRAMFS

EOF
chmod +x "$temp_dir/livekit/initramfs_create"

# Create initramfs image
print_step "Creating initramfs image"
pushd "$temp_dir/livekit"
initramfs_path=$(./initramfs_create "$kernel")
popd
cp "$initramfs_path" "initrfs-$kernel.img"
echo -e "  Initramfs image created: initrfs-$kernel.img"
